#include "connectionwidget.h"
#include "connectiondialog.h"
#include <DMessageBox>
#include <QVBoxLayout>
#include <QSettings>
#include <QSqlError>
#include <QSqlQuery>
#include <QMenu>

ConnectionWidget::ConnectionWidget(QWidget *parent)
    : DWidget(parent)
    , queryModel(new QSqlQueryModel(this))
{
    setupUI();
    setupConnections();
    loadConnections();
}

void ConnectionWidget::setupUI()
{
    auto layout = new QVBoxLayout(this);
    mainSplitter = new QSplitter(Qt::Horizontal, this);
    
    // Left panel
    auto leftWidget = new QWidget(mainSplitter);
    auto leftLayout = new QVBoxLayout(leftWidget);
    
    connectionsTree = new DTreeWidget(leftWidget);
    connectionsTree->setHeaderLabels({tr("Name"), tr("Status")});
    connectionsTree->setContextMenuPolicy(Qt::CustomContextMenu);
    
    auto btnLayout = new QHBoxLayout;
    addBtn = new DPushButton(tr("Add"), leftWidget);
    editBtn = new DPushButton(tr("Edit"), leftWidget);
    removeBtn = new DPushButton(tr("Remove"), leftWidget);
    testBtn = new DPushButton(tr("Test"), leftWidget);
    connectBtn = new DPushButton(tr("Connect"), leftWidget);
    disconnectBtn = new DPushButton(tr("Disconnect"), leftWidget);
    refreshBtn = new DPushButton(tr("Refresh"), leftWidget);
    
    btnLayout->addWidget(addBtn);
    btnLayout->addWidget(editBtn);
    btnLayout->addWidget(removeBtn);
    btnLayout->addWidget(testBtn);
    btnLayout->addWidget(connectBtn);
    btnLayout->addWidget(disconnectBtn);
    btnLayout->addWidget(refreshBtn);
    
    leftLayout->addWidget(connectionsTree);
    leftLayout->addLayout(btnLayout);
    
    // Right panel
    rightSplitter = new QSplitter(Qt::Vertical, mainSplitter);
    
    queryEditor = new DTextEdit(rightSplitter);
    resultsView = new DTableView(rightSplitter);
    
    auto queryWidget = new QWidget(rightSplitter);
    auto queryLayout = new QVBoxLayout(queryWidget);
    executeBtn = new DPushButton(tr("Execute"), queryWidget);
    
    queryLayout->addWidget(queryEditor);
    queryLayout->addWidget(executeBtn);
    
    resultsView->setModel(queryModel);
    
    mainSplitter->addWidget(leftWidget);
    mainSplitter->addWidget(rightSplitter);
    
    layout->addWidget(mainSplitter);
    
    // Set initial sizes
    mainSplitter->setSizes({200, 600});
    rightSplitter->setSizes({200, 400});
}

void ConnectionWidget::setupConnections()
{
    connect(addBtn, &DPushButton::clicked, this, &ConnectionWidget::addConnection);
    connect(editBtn, &DPushButton::clicked, this, &ConnectionWidget::editConnection);
    connect(removeBtn, &DPushButton::clicked, this, &ConnectionWidget::removeConnection);
    connect(testBtn, &DPushButton::clicked, this, &ConnectionWidget::testConnection);
    connect(connectBtn, &DPushButton::clicked, this, &ConnectionWidget::connectToDatabase);
    connect(disconnectBtn, &DPushButton::clicked, this, &ConnectionWidget::disconnectFromDatabase);
    connect(refreshBtn, &DPushButton::clicked, this, &ConnectionWidget::refreshDatabaseObjects);
    connect(executeBtn, &DPushButton::clicked, this, &ConnectionWidget::executeQuery);
    
    connect(connectionsTree, &DTreeWidget::customContextMenuRequested,
            this, &ConnectionWidget::showContextMenu);
    connect(connectionsTree, &DTreeWidget::itemDoubleClicked,
            this, &ConnectionWidget::handleTreeItemDoubleClicked);
}

void ConnectionWidget::addConnection()
{
    ConnectionDialog dialog(this);
    if (dialog.exec() == DDialog::Accepted) {
        QString name = dialog.connectionName();
        if (name.isEmpty()) {
            DMessageBox::warning(this, tr("Error"), 
                               tr("Connection name cannot be empty"));
            return;
        }
        
        if (connections.contains(name)) {
            DMessageBox::warning(this, tr("Error"), 
                               tr("Connection name already exists"));
            return;
        }
        
        QSqlDatabase db = QSqlDatabase::addDatabase(dialog.driverName(), name);
        db.setHostName(dialog.hostName());
        db.setPort(dialog.port());
        db.setDatabaseName(dialog.databaseName());
        db.setUserName(dialog.userName());
        db.setPassword(dialog.password());
        
        connections[name] = db;
        
        auto item = new QTreeWidgetItem(connectionsTree);
        item->setText(0, name);
        item->setText(1, tr("Disconnected"));
        item->setIcon(0, QIcon::fromTheme("database"));
        
        saveConnections();
    }
}

void ConnectionWidget::editConnection()
{
    auto item = connectionsTree->currentItem();
    if (!item)
        return;
        
    QString name = item->text(0);
    if (!connections.contains(name))
        return;
        
    ConnectionDialog dialog(this);
    dialog.setConnection(name, connections[name]);
    
    if (dialog.exec() == DDialog::Accepted) {
        // Remove old connection
        QSqlDatabase::removeDatabase(name);
        connections.remove(name);
        
        // Add new connection
        QString newName = dialog.connectionName();
        QSqlDatabase db = QSqlDatabase::addDatabase(dialog.driverName(), newName);
        db.setHostName(dialog.hostName());
        db.setPort(dialog.port());
        db.setDatabaseName(dialog.databaseName());
        db.setUserName(dialog.userName());
        db.setPassword(dialog.password());
        
        connections[newName] = db;
        
        item->setText(0, newName);
        item->setText(1, tr("Disconnected"));
        
        saveConnections();
    }
}

void ConnectionWidget::connectToDatabase()
{
    auto item = connectionsTree->currentItem();
    if (!item)
        return;
        
    QString name = item->text(0);
    if (!connections.contains(name))
        return;
        
    QSqlDatabase db = connections[name];
    if (db.open()) {
        updateConnectionStatus(name, true);
        loadDatabaseObjects(name);
    } else {
        DMessageBox::warning(this, tr("Error"), 
                           tr("Failed to connect: %1").arg(db.lastError().text()));
    }
}

void ConnectionWidget::disconnectFromDatabase()
{
    auto item = connectionsTree->currentItem();
    if (!item)
        return;
        
    QString name = item->text(0);
    if (!connections.contains(name))
        return;
        
    QSqlDatabase db = connections[name];
    db.close();
    updateConnectionStatus(name, false);
    
    // Clear database objects
    while (item->childCount() > 0) {
        delete item->takeChild(0);
    }
}

void ConnectionWidget::loadDatabaseObjects(const QString &connectionName)
{
    auto item = nullptr;
    for (int i = 0; i < connectionsTree->topLevelItemCount(); ++i) {
        if (connectionsTree->topLevelItem(i)->text(0) == connectionName) {
            item = connectionsTree->topLevelItem(i);
            break;
        }
    }
    
    if (!item)
        return;
        
    // Clear existing items
    while (item->childCount() > 0) {
        delete item->takeChild(0);
    }
    
    QSqlDatabase db = connections[connectionName];
    
    // Add Tables node
    auto tablesNode = new QTreeWidgetItem(item);
    tablesNode->setText(0, tr("Tables"));
    tablesNode->setIcon(0, QIcon::fromTheme("folder"));
    
    // Add tables
    auto tables = db.tables();
    for (const QString &table : tables) {
        addTableNode(tablesNode, table);
    }
    
    item->setExpanded(true);
}

void ConnectionWidget::addTableNode(QTreeWidgetItem *parent, const QString &tableName)
{
    auto tableItem = new QTreeWidgetItem(parent);
    tableItem->setText(0, tableName);
    tableItem->setIcon(0, QIcon::fromTheme("table"));
    
    // Add Columns node
    auto columnsNode = new QTreeWidgetItem(tableItem);
    columnsNode->setText(0, tr("Columns"));
    columnsNode->setIcon(0, QIcon::fromTheme("folder"));
    
    // TODO: Add columns information
}

void ConnectionWidget::executeQuery()
{
    auto item = connectionsTree->currentItem();
    if (!item)
        return;
        
    // Find top level item (connection)
    while (item->parent()) {
        item = item->parent();
    }
    
    QString connectionName = item->text(0);
    if (!connections.contains(connectionName))
        return;
        
    QSqlDatabase db = connections[connectionName];
    if (!db.isOpen()) {
        DMessageBox::warning(this, tr("Error"), 
                           tr("Database is not connected"));
        return;
    }
    
    QString query = queryEditor->toPlainText().trimmed();
    if (query.isEmpty()) {
        DMessageBox::warning(this, tr("Error"), 
                           tr("Query is empty"));
        return;
    }
    
    queryModel->setQuery(query, db);
    
    if (queryModel->lastError().isValid()) {
        DMessageBox::warning(this, tr("Error"), 
                           tr("Query failed: %1").arg(queryModel->lastError().text()));
    }
}

void ConnectionWidget::showContextMenu(const QPoint &pos)
{
    auto item = connectionsTree->itemAt(pos);
    if (!item)
        return;
        
    QMenu menu(this);
    
    // Find top level item (connection)
    auto connectionItem = item;
    while (connectionItem->parent()) {
        connectionItem = connectionItem->parent();
    }
    
    QString connectionName = connectionItem->text(0);
    bool isConnected = connections.contains(connectionName) && 
                      connections[connectionName].isOpen();
    
    // Add actions based on item type
    if (item == connectionItem) {
        // Connection item
        menu.addAction(tr("Edit"), this, &ConnectionWidget::editConnection);
        menu.addAction(tr("Remove"), this, &ConnectionWidget::removeConnection);
        menu.addSeparator();
        if (!isConnected)
            menu.addAction(tr("Connect"), this, &ConnectionWidget::connectToDatabase);
        else
            menu.addAction(tr("Disconnect"), this, &ConnectionWidget::disconnectFromDatabase);
        menu.addAction(tr("Refresh"), this, &ConnectionWidget::refreshDatabaseObjects);
    } else if (item->text(0) == tr("Tables")) {
        // Tables node
        if (isConnected)
            menu.addAction(tr("Refresh"), this, &ConnectionWidget::refreshDatabaseObjects);
    } else {
        // Table item
        if (isConnected) {
            menu.addAction(tr("Select All"), [this, item]() {
                QString table = item->text(0);
                queryEditor->setText(QString("SELECT * FROM %1;").arg(table));
                executeQuery();
            });
            // Add more table operations...
        }
    }
    
    if (!menu.isEmpty())
        menu.exec(connectionsTree->viewport()->mapToGlobal(pos));
}

void ConnectionWidget::handleTreeItemDoubleClicked(QTreeWidgetItem *item, int column)
{
    if (!item->parent()) {
        // Connection item
        if (item->text(1) == tr("Disconnected"))
            connectToDatabase();
        else
            disconnectFromDatabase();
    } else if (item->text(0) == tr("Tables")) {
        // Tables node
        item->setExpanded(!item->isExpanded());
    } else {
        // Table item
        QString table = item->text(0);
        queryEditor->setText(QString("SELECT * FROM %1;").arg(table));
        executeQuery();
    }
}

void ConnectionWidget::loadConnections()
{
    QSettings settings;
    settings.beginGroup("DatabaseConnections");
    
    auto connectionNames = settings.childGroups();
    for (const QString &name : connectionNames) {
        settings.beginGroup(name);
        
        QString driver = settings.value("driver").toString();
        QString host = settings.value("host").toString();
        int port = settings.value("port").toInt();
        QString database = settings.value("database").toString();
        QString user = settings.value("user").toString();
        QString password = settings.value("password").toString();
        
        QSqlDatabase db = QSqlDatabase::addDatabase(driver, name);
        db.setHostName(host);
        db.setPort(port);
        db.setDatabaseName(database);
        db.setUserName(user);
        db.setPassword(password);
        
        connections[name] = db;
        
        auto item = new QTreeWidgetItem(connectionsTree);
        item->setText(0, name);
        item->setText(1, tr("Disconnected"));
        item->setIcon(0, QIcon::fromTheme("database"));
        
        settings.endGroup();
    }
    
    settings.endGroup();
}

void ConnectionWidget::saveConnections()
{
    QSettings settings;
    settings.beginGroup("DatabaseConnections");
    
    // Clear existing connections
    settings.remove("");
    
    // Save current connections
    for (auto it = connections.constBegin(); it != connections.constEnd(); ++it) {
        const QString &name = it.key();
        const QSqlDatabase &db = it.value();
        
        settings.beginGroup(name);
        settings.setValue("driver", db.driverName());
        settings.setValue("host", db.hostName());
        settings.setValue("port", db.port());
        settings.setValue("database", db.databaseName());
        settings.setValue("user", db.userName());
        settings.setValue("password", db.password());
        settings.endGroup();
    }
    
    settings.endGroup();
}

void ConnectionWidget::updateConnectionStatus(const QString &name, bool connected)
{
    for (int i = 0; i < connectionsTree->topLevelItemCount(); ++i) {
        auto item = connectionsTree->topLevelItem(i);
        if (item->text(0) == name) {
            item->setText(1, connected ? tr("Connected") : tr("Disconnected"));
            break;
        }
    }
}

void ConnectionWidget::refreshDatabaseObjects()
{
    auto item = connectionsTree->currentItem();
    if (!item)
        return;
        
    // Find top level item (connection)
    while (item->parent()) {
        item = item->parent();
    }
    
    QString name = item->text(0);
    if (connections.contains(name) && connections[name].isOpen()) {
        loadDatabaseObjects(name);
    }
}